Large-Scale Discrete Fourier Transform on TPUs

In this work, we present two parallel algorithms for the large-scale discrete Fourier transform (DFT) on Tensor Processing Unit (TPU) clusters. The are associated with DFT formulations: one formulation, denoted as KDFT, is based Kronecker product; other famous Cooley-Tukey algorithm and phase adjustment, FFT. Both KDFT FFT formulations take full advantage of TPU's strength in matrix multiplications. formulation allows direct use nonuniform inputs without additional step. algorithms, same strategy data decomposition applied to input data. Through decomposition, dense multiplications kept local within TPU cores, which can be performed completely parallel. communication among cores achieved through one-shuffle scheme both sending receiving takes place simultaneously between neighboring along direction interconnect network. designed topology clusters, minimizing time required by cores. implemented TensorFlow. three-dimensional complex an example dimension 8192 ×8192 a Pod: run 12.66 seconds that 8.3 seconds. Scaling analysis provided demonstrate high efficiency implementations TPUs.